import torch
import torch.nn as nn
from mmcls.models import BACKBONES
from mmcls.models.backbones import VisionTransformer
from mmcls.models.utils import resize_pos_embed
from typing import List


@BACKBONES.register_module()
class SVDIMT(VisionTransformer):

    def __init__(self,
                 prompt_length: int = 1,
                 prompt_layers: List[int] = None,
                 prompt_pos: str = 'postpend',
                 prompt_init: str = 'normal',
                 *args,
                 **kwargs):
        super().__init__(*args, **kwargs)
        for param in self.parameters():
            param.requires_grad = False


        # self.prompt_layers = [0] if prompt_layers is None else prompt_layers  
        self.prompt_layers = range(12) if prompt_layers is None else prompt_layers
        # n_patches = int((self.img_size[0]/16)*(self.img_size[1]/16))
        prompt = torch.empty(len(self.prompt_layers), 577, prompt_length)
        
        # for layer in self.prompt_layers:
        #     # linear_layer = nn.Linear(self.embed_dims+1, self.embed_dims*3,bias=True)# qkv:(769,768*3)
        #     linear_layer = nn.Linear(self.embed_dims+prompt_length, self.embed_dims)
        #     self.add_module(f'linear_{layer}', linear_layer)


        if prompt_init == 'uniform':
            nn.init.uniform_(prompt, -0.08, 0.08)
        elif prompt_init == 'zero':
            nn.init.zeros_(prompt)
        elif prompt_init == 'kaiming':
            nn.init.kaiming_normal_(prompt)
        elif prompt_init == 'token':
            nn.init.zeros_(prompt)
            self.prompt_initialized = False
        else:
            nn.init.normal_(prompt, std=0.02)
        self.prompt = nn.Parameter(prompt, requires_grad=True)
        self.prompt_length = prompt_length
        self.prompt_pos = prompt_pos


    def forward(self, x):
        """Following mmcls implementation."""
        B = x.shape[0] # batch size
        x, patch_resolution = self.patch_embed(x)

        # stole cls_tokens impl from Phil Wang, thanks
        cls_tokens = self.cls_token.expand(B, -1, -1)
        x = torch.cat((cls_tokens, x), dim=1) # x: (B,577,768)
        x = x + resize_pos_embed(
            self.pos_embed,
            self.patch_resolution,
            patch_resolution,
            mode=self.interpolate_mode,
            num_extra_tokens=self.num_extra_tokens) # x: (B,577,768)
        x = self.drop_after_pos(x) # x: (B,577,768)

        # self.prompt: (1=layer,577,1)
        # prompt: (1=layer,4,577,1)
        prompt = self.prompt.unsqueeze(1).expand(-1, x.shape[0], -1, -1)

        outs = []
        for i, layer in enumerate(self.layers):
            if i in self.prompt_layers:
                if self.prompt_pos == 'postpend':
                    # prompt: [layer, batch, length, dim]->(1,4,577,1)
                    ## x: (B,577,768) -> x: (B, 577,769)
                    x_769 = torch.cat([x[:, :, :], prompt[i, :, :, :]], dim=2)
    
            # 重写 x = layer(x)
                ## 得到attn_x
                def svd(x_):
                    x = x_[:, :, :self.embed_dims]+x_[:, :, self.embed_dims:].mean(dim=-1,keepdim=True)
                    # x = x_[:, :, :self.embed_dims]+x_[:, :, self.embed_dims:].mean(dim=-1,keepdim=True)
                    # x = linear(x_)
                    return x
                # 目的是实现Norm
                svd_x = svd(x_769)
                shotcut = False
                if shotcut:
                    svd_x = x+svd_x

                x = layer(svd_x)
            else:
                x = layer(x)

            if i == len(self.layers) - 1 and self.final_norm:
                x = self.norm1(x)

            if i in self.out_indices:
                outs.append(x[:, 0])

        return tuple(outs)